package io.renren.membermanage.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.renren.membermanage.dto.ModelTrainParamDto;
import io.renren.membermanage.dto.SseCommonMessageDto;
import io.renren.membermanage.entity.ModelTrain;
import io.renren.membermanage.entity.ModelTrainDetail;
import io.renren.membermanage.service.*;
import io.renren.sdapi.service.SseService;
import io.renren.sdapi.service.StableDiffusionApiService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Slf4j
@Service
public class TaskServiceImpl implements TaskService {

    @Resource
    private StableDiffusionApiService stableDiffusionApiService;

    @Resource
    private SseService service;

    @Resource
    private ModelTrainService modelTrainService;

    @Resource
    private ModelTrainDetailService modelTrainDetailService;

    @Resource
    private ModelTrainParamService modelTrainParamService;

    @Value("${train.log}")
    private String trainLog;

    @Value("${train.file.check}")
    private String trainFileCheck;

    @Value("${model.train.preview}")
    private String previewPath;

    private final ObjectMapper objectMapper = new ObjectMapper();

    public Runnable createTask(ModelTrain data){
        return new Task(data, stableDiffusionApiService, service, modelTrainService, modelTrainDetailService, modelTrainParamService);
    }

    private class Task implements Runnable {

        private final ModelTrain data;
        private final StableDiffusionApiService stableDiffusionApiService;
        private final SseService service;
        private final ModelTrainService modelTrainService;
        private final ModelTrainDetailService modelTrainDetailService;
        private final ModelTrainParamService modelTrainParamService;

        public Task(ModelTrain data, StableDiffusionApiService stableDiffusionApiService, SseService service,
                    ModelTrainService modelTrainService, ModelTrainDetailService modelTrainDetailService, ModelTrainParamService modelTrainParamService){
            this.data = data;
            this.stableDiffusionApiService = stableDiffusionApiService;
            this.service = service;
            this.modelTrainService = modelTrainService;
            this.modelTrainDetailService = modelTrainDetailService;
            this.modelTrainParamService = modelTrainParamService;
        }

        @Override
        public void run() {
            String workDic = System.getProperty("user.dir");
            workDic = workDic.substring(0,workDic.length()-1);
            workDic = workDic.substring(0,workDic.lastIndexOf("\\"));
            log.info("--------------Discover that the task has started training, and identify the primary key of the current training：{}", data.getId());
            ModelTrainParamDto param = null;
            try {
                param = objectMapper.readValue(data.getTrainParam(), ModelTrainParamDto.class);
            } catch (JsonProcessingException e) {
                throw new RuntimeException(e);
            }

            String taskId = executeTask(data.getId().toString(), param);
            if (StrUtil.isBlank(taskId)){
                log.error("There was a problem calling the interface for model training!");
                data.setStatus(9);
                data.setTrainTime(LocalDateTime.now());
                data.setUpdateTime(LocalDateTime.now());
                modelTrainService.updateById(data);
                return;
            }
            new Thread(() -> {
                data.setStatus(2);
                data.setTaskId(taskId);
                data.setTrainTime(LocalDateTime.now());
                data.setUpdateTime(LocalDateTime.now());
                modelTrainService.updateById(data);
            }).start();
            String fileName = param.getOutput_name() + ".safetensors";
            while (!isFileGenerated(workDic + trainFileCheck, fileName)) {
                try {
                    log.info("------------- Current user : " + data.getUserId() + "Be performing a task! Time with : " + LocalDateTime.now() + "--------------");
                    TimeUnit.SECONDS.sleep(60);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
            try {
                log.info("------------- Current user : " + data.getUserId() + " Delay waiting for the file to be fully generated" + LocalDateTime.now() + "--------------");
                TimeUnit.SECONDS.sleep(60);
            } catch (InterruptedException e) {
                return;
            }
            boolean flag = waitForFileGeneration(taskId, data.getId());
            if (flag){
                handleData(workDic + trainFileCheck, fileName, data.getUserId(), data.getId(), workDic);
                log.info("-----------------------The current training is successful. The primary key identifier of the training:{}", data.getId());
            }else {
                log.info("-----------------------The current training is failed. The primary key identifier of the training:{}", data.getId());
            }
        }

        private String executeTask(String modelTrainId, ModelTrainParamDto dto){
            return stableDiffusionApiService.modelTrain(modelTrainId, dto);
        }

        /**
         * 等待文件生成
         * @param taskId
         * @param modelTrainId
         * @return
         */
        private boolean waitForFileGeneration(String taskId, Long modelTrainId) {
            String workDic = System.getProperty("user.dir");
            workDic = workDic.substring(0,workDic.length()-1);
            workDic = workDic.substring(0,workDic.lastIndexOf("\\"));
            String taskCreatedPattern = "Task " + taskId + " created";
            Path logFilePath = Paths.get(workDic + trainLog);
            String failedFlag = "Training failed";
            String successFlag = "Training finished";
            try (BufferedReader reader = Files.newBufferedReader(logFilePath, StandardCharsets.UTF_8)) {
                String line;
                boolean identifierFound = false;
                while ((line = reader.readLine()) != null) {
                    if (line.contains(taskCreatedPattern)) {
                        identifierFound = true;
                    }
                    if (identifierFound){
                        if (line.contains(successFlag)){
                            return true;
                        }
                        if (line.contains(failedFlag)) {
                            LambdaQueryWrapper<ModelTrain> wrapper = new LambdaQueryWrapper<>();
                            wrapper.eq(ModelTrain::getId, modelTrainId);
                            ModelTrain update = modelTrainService.getOne(wrapper);
                            update.setStatus(9);
                            update.setUpdateTime(LocalDateTime.now());
                            modelTrainService.updateById(update);
                            return false;
                        }
                    }
                }
            }catch (IOException e) {
                e.printStackTrace();
            }
            return false;
        }

        private boolean isFileGenerated(String trainFileCheck, String fileName) {
            File folder = new File(trainFileCheck + fileName);
            return folder.exists();
        }

        private void handleData(String trainFileCheck, String fileName, Long userId, Long modelTrainId, String workDic){
            if (isFileGenerated(trainFileCheck, fileName)) {
                log.info("------------- Current user : " + userId + "File generated! Path is : " + trainFileCheck + fileName + "------------");
                //训练完成，将模型名和预览图插入到库开始
                String regex = fileName.replace(".safetensors","")+"-.*\\.safetensors"; // 模糊匹配模式
                File directory = new File(trainFileCheck);
                List<String> modelFileName = new ArrayList<String>();
                List<String> previeFileName = new ArrayList<String>();
                if (directory.isDirectory()) {
                    File[] files = directory.listFiles();
                    if (files != null) {
                        for (File file : files) {
                            if (file.isFile()) {
                                String modelName = file.getName();
                                Pattern pattern = Pattern.compile(regex);
                                Matcher matcher = pattern.matcher(modelName);
                                if (matcher.matches()) {
                                    modelFileName.add(modelName);
                                }
                            }
                        }
                    }
                } else {
                    System.out.println("指定路径不是一个目录。");
                }
                directory = new File(trainFileCheck + "sample/");
                regex = fileName.replace(".safetensors","")+"\\_.*\\.png"; // 模糊匹配模式
                if (directory.isDirectory()) {
                    File[] files = directory.listFiles();
                    if (files != null) {
                        for (File file : files) {
                            if (file.isFile()) {
                                String previeeName = file.getName();
                                Pattern pattern = Pattern.compile(regex);
                                Matcher matcher = pattern.matcher(previeeName);
                                if (matcher.matches()) {
                                    previeFileName.add(previeeName);
                                }
                            }
                        }
                    }
                } else {
                    System.out.println("指定路径不是一个目录。");
                }
                log.info(JSONUtil.toJsonStr(modelFileName)+"----------文件名--------"+JSONUtil.toJsonStr(previeFileName));
                List<ModelTrainDetail> modelTrainDetails = new ArrayList<>();
                ModelTrain byId = modelTrainService.getById(modelTrainId);
                for(int i=0;i<modelFileName.size();i++){
                    String oneModelName = modelFileName.get(i).
                            replace(".safetensors","")
                            .replace("-","_e");
                    for(int j=0;j<previeFileName.size();j++){
                        if(previeFileName.get(j).contains(oneModelName)){
                            ModelTrainDetail modelTrainDetail = new ModelTrainDetail();
                            modelTrainDetail.setModelTrainId(modelTrainId);
                            modelTrainDetail.setModelFileName(modelFileName.get(i));
                            modelTrainDetail.setPreviewFileName(previeFileName.get(j));
                            modelTrainDetail.setStep(modelTrainDetail.getModelFileName()
                                    .replace(byId.getOutputName(), "")
                                    .replace("-","")
                                    .replace(".safetensors",""));
                            modelTrainDetails.add(modelTrainDetail);
                            previeFileName.remove(j);
                            //新增（非必要代码）复制当前的图片到nginx目录下
                            String currentPath = directory.getPath() + previeFileName.get(j);
                            //源文件路径
                            Path sourcePath = Paths.get(currentPath);
                            // 目标文件夹路径
                            Path targetDirectory = Paths.get(workDic + previewPath);
                            Path targetPath = targetDirectory.resolve(sourcePath.getFileName());
                            try {
                                // 复制文件
                                Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
                                System.out.println("File copy success！");
                            } catch (IOException e) {
                                System.err.println("File copy failure: " + e.getMessage());
                            }
                            break;
                        }
                    }
                }
                log.info(JSONUtil.toJsonStr(modelTrainDetails));
                ModelTrainDetail modelTrainDetail = new ModelTrainDetail();
                modelTrainDetail.setModelTrainId(modelTrainId);
                modelTrainDetail.setModelFileName(fileName);
                modelTrainDetail.setPreviewFileName(previeFileName.get(0));
                String stepStr = modelTrainDetail.getPreviewFileName().substring(modelTrainDetail.getPreviewFileName().indexOf("_e"));
                stepStr = stepStr.substring(0, stepStr.indexOf("_"));
                modelTrainDetail.setStep(stepStr);
                modelTrainDetails.add(modelTrainDetail);
                modelTrainDetailService.saveBatch(modelTrainDetails);
                //训练完成，将模型名和预览图插入到库结束

                SseCommonMessageDto sseMessageDto = new SseCommonMessageDto();
                sseMessageDto.setUserId(userId.toString());
                sseMessageDto.setData("The training has been completed. Please go and check the file path：" + trainFileCheck + fileName);
                sseMessageDto.setTotal(1);
                service.sendMessage(sseMessageDto);
                byId.setStatus(1);
                byId.setGenerateAddress(trainFileCheck + fileName);
                byId.setUpdateTime(LocalDateTime.now());
                byId.setTrainedTime(LocalDateTime.now());
                modelTrainService.updateById(byId);
            }
        }
    }
}
